{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imaging with a vortex coronagraph\n",
    "\n",
    "We will simulate on-axis and off-axis images of stars through a (ring-apodized) vortex coronagraph.\n",
    "\n",
    "We'll start by importing all relevant libraries and setting up our pupil and focal grids. We'll slightly oversize our pupil grid to more clearly see the effects of the vortex coronagraph in the Lyot plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from hcipy import *\n",
    "import numpy as np\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# For notebook animations\n",
    "from matplotlib import animation\n",
    "from IPython.display import HTML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rcParams['figure.dpi'] = 100\n",
    "\n",
    "pupil_grid = make_pupil_grid(256, 1.5)\n",
    "focal_grid = make_focal_grid(8, 12)\n",
    "prop = FraunhoferPropagator(pupil_grid, focal_grid)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start of by creating a circular aperture. A vortex coronagraph works perfectly for a circular aperture. We'll use supersampling to evaluate this aperture to partially suppress sampling artefacts. We'll also use a slightly-undersized circular Lyot stop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aperture = evaluate_supersampled(circular_aperture(1), pupil_grid, 4)\n",
    "lyot_mask = evaluate_supersampled(circular_aperture(0.95), pupil_grid, 4)\n",
    "\n",
    "plt.subplot(1,2,1)\n",
    "plt.title('Aperture')\n",
    "imshow_field(aperture, cmap='gray')\n",
    "plt.subplot(1,2,2)\n",
    "plt.title('Lyot stop')\n",
    "imshow_field(lyot_mask, cmap='gray')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can perform non-coronagraphic imaging by just using the Fraunhofer propagation defined above, to propagate the light from the pupil to the focal plane of our telescope."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wf = Wavefront(aperture)\n",
    "img_ref = prop(wf).intensity\n",
    "\n",
    "imshow_field(np.log10(img_ref / img_ref.max()), vmin=-5, cmap='inferno')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This shows the usual Airy pattern. We'll now generate the vortex coronagraph itself. It requires a pupil grid and the charge of the vortex. The vortex coronagraph object propagates light from the pupil plane to the Lyot plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "charge = 2\n",
    "coro = VortexCoronagraph(pupil_grid, charge)\n",
    "lyot_stop = Apodizer(lyot_mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now propagate light through the vortex coronagraph. Internally the vortex coronagraph performs many propagations through the vortex with successively higher resolutions. This is done to adequately sample the vortex singularity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wf = Wavefront(aperture)\n",
    "lyot_plane = coro(wf)\n",
    "\n",
    "imshow_field(lyot_plane.intensity, cmap='inferno')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now block this light with a Lyot stop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "post_lyot_mask = lyot_stop(lyot_plane)\n",
    "img = prop(post_lyot_mask).intensity\n",
    "\n",
    "imshow_field(np.log10(img / img_ref.max()), vmin=-5, vmax=0, cmap='inferno')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The star has completely been suppressed. We can see the star appear again, when we look at an off-axis object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wf = Wavefront(aperture * np.exp(2j * np.pi * pupil_grid.x * 1.5))\n",
    "img = prop(lyot_stop(coro(wf))).intensity\n",
    "\n",
    "imshow_field(np.log10(img / img_ref.max()), vmin=-5, vmax=0, cmap='inferno')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the Lyot plane image looks totally different for an off-axis star:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lyot = coro(wf)\n",
    "\n",
    "imshow_field(lyot.intensity, vmax=2, cmap='inferno')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unintuitively, the light in the Lyot stop is offset in the vertical direction, while the star is offset in the horizontal direction. We can see this effect clearer in an animation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_offaxis_animation(coro):\n",
    "    fig = plt.figure()\n",
    "    plt.subplot(1,2,1)\n",
    "    plt.title('Lyot plane')\n",
    "    im1 = imshow_field(lyot_plane.intensity, vmax=1, cmap='inferno')\n",
    "    plt.subplot(1,2,2)\n",
    "    plt.title('Science image plane')\n",
    "    im2 = imshow_field(np.log10(img / img_ref.max()), vmin=-5, vmax=0, cmap='inferno')\n",
    "    plt.close(fig)\n",
    "\n",
    "    def animate(angular_separation):\n",
    "        wf = Wavefront(aperture * np.exp(2j * np.pi * pupil_grid.x * angular_separation))\n",
    "        lyot = coro(wf)\n",
    "        img = prop(lyot_stop(lyot))\n",
    "\n",
    "        im1.set_data(*pupil_grid.separated_coords, lyot.intensity.shaped)\n",
    "        im2.set_data(*focal_grid.separated_coords, np.log10(img.intensity.shaped / img_ref.max()))\n",
    "\n",
    "        return [im1, im2]\n",
    "\n",
    "    angular_separations = np.linspace(-5, 5, 51)\n",
    "    anim = animation.FuncAnimation(fig, animate, angular_separations, interval=160, blit=True)\n",
    "    return HTML(anim.to_html5_video())\n",
    "\n",
    "create_offaxis_animation(coro)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also simulate vortex coronagraphs with other charges:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vortex4 = VortexCoronagraph(pupil_grid, charge=4)\n",
    "create_offaxis_animation(vortex4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  },
  "level": "beginner",
  "thumbnail_figure_index": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
